from math import sin, cos, pi
from stl import mesh
import numpy as np
import os
import tqdm
from scipy.interpolate import RectBivariateSpline
import cv2

"""
适用于通过简单圆柱投影的DEM数据，以生成球面三维模型stl文件。
参数说明：
    X_PIXEL_RANGE (int) : x方向的像素范围
    Y_PIXEL_RANGE (int) : y方向的像素范围
    LATITUDE_RANGE (float) : 纬度范围，单位为弧度
    LONGITUDE_RANGE (float) : 经度范围，单位为弧度
    base_batch (int) : 每个子图块的像素尺寸（以正方形）
    rate (int) : 降采样率，即在原图上每隔几个像素取一个像素
    raw_reso (float) : DEM在赤道上的分辨率（即未失真的分辨率），单位为米/pix
    tags (str) : 用于区分不同的球面模型，命名参数
    theta_start (float) : 球面模型的左上角的纬度，单位为弧度
    phi_start (float) : 球面模型的左上角的经度，单位为弧度
    path (str) : DEM数据的路径
    output_dir (str) : 输出文件的路径
"""
X_PIXEL_RANGE = 48131
Y_PIXEL_RANGE = 8264
LATITUDE_RANGE = (45.2410 - 42.8613) / 180 * pi
LONGITUDE_RANGE = (-51.6879 - (-52.0974)) / 180 * pi
base_batch = 1000
rate = 2
raw_reso = 1.5
tags = "sp"  # sphere
theta_start = 45.2410 / 180 * pi
phi_start = -52.0974 / 180 * pi
path = (
    "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/regularized_M1173414625_DEM.tif"
)
output_dir = "/disk527/sdb1/a804_cbf/datasets/lunar_crater"


def sphere2plane(lat, lon, lat0=45.2410 / 180 * pi, lon0=-52.0974 / 180 * pi):
    """
    Arguments:
        lat (float) : latitude in degree
        lon (float) : longitude in degree
        lat0 (float) : latitude in degree
        lon0 (float) : longitude in degree
    Return:
        x (float) : x in pixel
        y (float) : y in pixel
    """
    x = (lat0 - lat) * X_PIXEL_RANGE / LATITUDE_RANGE
    y = (lon - lon0) * Y_PIXEL_RANGE / LONGITUDE_RANGE
    return x, y


# 使用极坐标系解决问题是可行的
def create_sphere_segment(
    dem_image,
    row,
    col,
    radius,
    theta_start,
    delta_theta,
    phi_start,
    delta_phi,
    output_path,
    rate=2,
):
    # 生成球面片段的顶点
    vertices = []
    x = np.arange(0, dem_image.shape[0])
    y = np.arange(0, dem_image.shape[1])
    f = RectBivariateSpline(x, y, dem_image / raw_reso)
    num_steps = dem_image.shape[0] // rate

    for i in range(num_steps):
        for j in range(num_steps):
            # 纬度逐渐减小，经度逐渐增加
            theta = theta_start - i / (num_steps - 1) * delta_theta
            phi = phi_start + j / (num_steps - 1) * delta_phi
            xnew, ynew = sphere2plane(theta, phi)
            # 就是这里造成了接缝
            r = (
                f(
                    xnew - row * (dem_image.shape[0] - rate),
                    ynew - col * (dem_image.shape[1] - rate),
                )[0, 0]
                + radius / raw_reso
            ) / rate
            assert (
                xnew - row * (dem_image.shape[0] - rate)
                <= dem_image.shape[0] + rate + 1
                and xnew - row * (dem_image.shape[0] - rate) >= -rate - 1
            ) and (
                ynew - col * (dem_image.shape[1] - rate)
                <= dem_image.shape[1] + rate + 1
                and ynew - col * (dem_image.shape[1] - rate) >= -rate - 1
            ), f"Out of bound, {xnew -row * (dem_image.shape[0]-rate)}, {ynew - col * (dem_image.shape[1]-rate)}"
            x = r * cos(theta) * cos(phi)
            y = r * cos(theta) * sin(phi)
            z = r * sin(theta)
            vertices.append((x, y, z))

    # 生成三角形
    faces = []
    for i in range(num_steps - 1):
        for j in range(num_steps - 1):
            # 当前点的索引
            current = i * num_steps + j
            # 下一个theta值的点的索引
            next_theta = current + num_steps
            faces.append([current, next_theta, current + 1])
            faces.append([current + 1, next_theta, next_theta + 1])

    # 创建网格
    sphere_segment = mesh.Mesh(np.zeros(len(faces), dtype=mesh.Mesh.dtype))
    for i, f in enumerate(faces):
        for j in range(3):
            sphere_segment.vectors[i][j] = vertices[f[j]]

        # 保存为STL文件
    sphere_segment.save(output_path)


if __name__ == "__main__":

    num_batches = X_PIXEL_RANGE // base_batch
    num_samples = Y_PIXEL_RANGE // base_batch
    # 参数
    delta_theta = LATITUDE_RANGE / X_PIXEL_RANGE * base_batch
    delta_phi = LONGITUDE_RANGE / Y_PIXEL_RANGE * base_batch
    radius = 1737.4 * 1000
    print(
        f"The basic batch of DEM is {base_batch}x{base_batch}. Each line contains {num_samples} batches. There are totally {num_batches*num_samples} basic batches."
    )
    image = cv2.imread(path, cv2.IMREAD_UNCHANGED)

    for row in tqdm.tqdm(range(num_batches)):
        if not os.path.exists(os.path.join(output_dir, "stl_models")):
            os.makedirs(os.path.join(output_dir, "stl_models"))
        if not os.path.exists(os.path.join(output_dir, "stl_models", f"{row}")):
            os.makedirs(os.path.join(output_dir, "stl_models", f"{row}"))
        for col in range(num_samples):
            img_array = image[
                row * base_batch : (row + 1) * base_batch + rate,
                col * base_batch : (col + 1) * base_batch + rate,
            ]
            # 创建球面片段
            num_steps = base_batch // rate
            create_sphere_segment(
                img_array,
                row,
                col,
                radius,
                theta_start - delta_theta * row,  # 北半球
                delta_theta * (base_batch + rate) / base_batch,
                phi_start + delta_phi * col,  # 东半球
                delta_phi * (base_batch + rate) / base_batch,
                os.path.join(
                    output_dir,
                    "stl_models",
                    f"{row}",
                    f"chang_e_{tags}_{row}_{col}_{num_steps}_{raw_reso}m.stl",
                ),
                rate,
            )
            np.savez(
                os.path.join(
                    output_dir,
                    "dem_int16",
                    f"{row}",
                    f"chang_e_{row}_{col}_{base_batch}.npz",
                ),
                image=img_array,
            )
